"""
分散任务执行时间
"""
import copy
import math
import datetime
import arrow
import requests
from app.config.settings import SCANNER_SERVER
from app.db.models.wass import ScanEvents
from app.db.models.users import User


client = ScanEvents._get_db()


class DecentralizeTask(object):

    def __init__(self, is_only_knownsec=True, is_only_scanv=True):
        # 是否只处理ScanV系统公司账号用户的资产
        self.is_only_knownsec = is_only_knownsec
        # 是否只处理ScanV系统资产
        self.is_only_scanv = is_only_scanv
        self.collection = client["scan_tasks"]
        self.collection_bk = client["tmp_task_data"]
        self.task_config_map = {}
        self.task_data_map = {}
        self.get_all_task_time()

    @staticmethod
    def _parse_periodic_time(conf):
        """将周期全部换算成秒"""
        seconds = conf.get("seconds", 0)
        if weeks := conf.get("weeks"):
            seconds += weeks * 7 * 24 * 3600
        if days := conf.get("days", 0):
            seconds += days * 24 * 3600
        if hours := conf.get("hours", 0):
            seconds += hours * 3600
        if minutes := conf.get("minutes", 0):
            seconds += minutes * 60
        return seconds

    def get_all_task_time(self):
        query = {"task_config.is_periodic": True, "task_config.is_disabled": False,
                 "addition.taskType": {"$nin": ["http", "ping"]}}
        if self.is_only_scanv:
            query.update({"$or": [{"addition.notificationTarget": {"$exists": 0}},
                                  {"addition.notificationTarget": 1}]})
        if self.is_only_knownsec:
            uids = [str(u.id) for u in User.objects.find({"username": {"$regex": "@knownsec.com"}}).only("id")]
            query.update({"addition.uid": {"$in": uids}})
        for task in self.collection.find(query):
            periodic_time = self._parse_periodic_time(task.get("task_config", {}))
            # 过滤掉周期小于一天的任务
            if periodic_time < 86400:
                continue
            task_id = task.get("task_id")
            next_start_at = task.get("next_start_at")
            last_started_at = task.get("last_started_at")
            if (not next_start_at) and (not last_started_at):
                continue
            if not next_start_at:
                next_start_at = last_started_at + datetime.timedelta(seconds=periodic_time)
            task_conf = {"periodic_time": periodic_time, "next_start_at": next_start_at, "task_id": task_id}
            self.task_config_map.update({task_id: task_conf})
            self.task_data_map.update({task_id: {
                "task_id": task_id,
                "target_url": task["target_url"],
                "origin_ip": task["origin_ip"],
                "job_names": task["job_names"],
                "filter_events": task["filter_events"],
                "task_config": task["task_config"],
                "job_config_map": task["job_config_map"],
                "flex_job_config_map": task["flex_job_config_map"],
                "target_id": task["target_id"],
                "user_id": task["user_id"],
                "addition": task["addition"],
            }})

    def _add_second(self, seconds, task_list, time_step, task_step, start_second):
        """更新每小时时间"""
        now = datetime.datetime.utcnow()
        night_time = now.replace(hour=0, minute=0, second=0, microsecond=0)
        index = 0
        for i in range(0, seconds, time_step):
            _task_list = task_list[index * task_step: (index + 1) * task_step]
            index += 1
            for task in _task_list:
                new_nt = night_time + datetime.timedelta(seconds=start_second + i)
                old_nt = self.task_config_map[task["task_id"]]["next_start_at"]
                res_nt = old_nt.replace(hour=new_nt.hour, minute=new_nt.minute, second=new_nt.second, microsecond=new_nt.microsecond)
                task["next_start_at"] = res_nt
                self.task_config_map[task["task_id"]]["next_start_at"] = res_nt

    def _update_per_day(self, source_tasks):
        """计算一天中的任务分布"""
        task_list = list(copy.deepcopy(source_tasks))
        # 目前每天执行的任务数量还远远小于 24 * 3600， 就暂时不考虑大于24 * 3600的情况
        if len(source_tasks) > 24 * 3600:
            raise ValueError("任务太多")
        now = datetime.datetime.utcnow()
        for task in task_list:
            task["next_start_at"] = task["next_start_at"].replace(year=now.year, month=now.month, day=now.day)
        task_count = len(task_list)
        # 排序， 尽量保证先执行的任务还是先执行
        task_list = list(sorted(task_list, key=lambda x: x["next_start_at"]))
        if task_count <= 24:
            per_hour_task_count = 1
            second_step = 3600
        else:
            per_hour_task_count = math.ceil(task_count / 24)
            second_step = int(3600 / per_hour_task_count)
        for hour in range(24):
            _task = task_list[hour * per_hour_task_count: (hour + 1) * per_hour_task_count]
            self._add_second(3600, _task, second_step, 1, 3600 * hour)
        return task_list

    def _update_time(self):
        """更新所有任务在一天中的运行时间"""
        task_list = list(copy.deepcopy(self.task_config_map).values())
        task_periodic_map = {}
        seconds = 24 * 3600
        for task in task_list:
            periodic_time = task["periodic_time"]
            if not task_periodic_map.get(periodic_time):
                task_periodic_map[periodic_time] = [task]
            else:
                task_periodic_map[periodic_time].append(task)
        # 得出每天需要运行的任务
        per_day_tasks = []
        periodic_task_count_map = {}
        for periodic_time, periodic_tasks in task_periodic_map.items():
            task_count = len(periodic_tasks)
            if periodic_time <= seconds:
                per_day_tasks.extend(periodic_tasks)
                periodic_task_count_map.update({periodic_time: {"per_count": task_count, "periodic_task_count": task_count}})
            else:
                days = int(periodic_time / seconds)
                if task_count <= days:
                    per_day_tasks.extend(periodic_tasks[:1])
                    periodic_task_count_map.update({periodic_time: {"per_count": 1, "periodic_task_count": task_count}})
                else:
                    # 取值可能比每天实际数量大1
                    per_day_count = math.ceil(task_count / days)
                    per_day_tasks.extend(periodic_tasks[:per_day_count])
                    periodic_task_count_map.update({periodic_time: {"per_count": per_day_count, "periodic_task_count": task_count}})
        # 根据每天的任务来分布一天中的执行时间
        res_tasks = self._update_per_day(per_day_tasks)
        res_task_periodic_map = {}
        for task in res_tasks:
            periodic_time = task["periodic_time"]
            if not res_task_periodic_map.get(periodic_time):
                res_task_periodic_map[periodic_time] = [task]
            else:
                res_task_periodic_map[periodic_time].append(task)
        # 对周期大于一天的任务的运行时间需要根据上边每天的分布结果进行从新分配
        for periodic_time, periodic_tasks in task_periodic_map.items():
            if periodic_time > seconds:
                periodic_task_count = periodic_task_count_map[periodic_time]["periodic_task_count"]
                per_count = periodic_task_count_map[periodic_time]["per_count"]
                res_periodic_tasks = res_task_periodic_map[periodic_time]
                # 对每天的任务按小时排序， 这一步之前必须保证已经进行过按天分布
                periodic_tasks = list(sorted(periodic_tasks, key=lambda x: x["next_start_at"]))
                for i in range(0, math.ceil(periodic_task_count / per_count)):
                    _task_list = periodic_tasks[i * per_count: (i+1) * per_count]
                    # 在同一周期不同天执行的任务运行时间设置为一样
                    for index, task in enumerate(_task_list):
                        res_task = res_periodic_tasks[index]
                        new_nt = res_task["next_start_at"]
                        old_nt = self.task_config_map[task["task_id"]]["next_start_at"]
                        self.task_config_map[task["task_id"]]["next_start_at"] = \
                            old_nt.replace(hour=new_nt.hour, minute=new_nt.minute, second=new_nt.second, microsecond=new_nt.microsecond)

    def _add_date(self, task, days):
        now = datetime.datetime.utcnow()
        old_nt = self.task_config_map[task["task_id"]]["next_start_at"]
        new_nt = now + datetime.timedelta(days=days)
        self.task_config_map[task["task_id"]]["next_start_at"] = \
            new_nt.replace(hour=old_nt.hour, minute=old_nt.minute, second=old_nt.second, microsecond=old_nt.microsecond)

    def _update_day(self):
        """计算任务在一个周期内执行的日期"""
        # 将任务排序， 尽量减少偏移量
        task_list = sorted(copy.deepcopy(self.task_config_map).values(), key=lambda x: x["next_start_at"])
        task_periodic_map = {}
        # 将任务按照执行周期分类
        for task in task_list:
            periodic_time = task["periodic_time"]
            if not task_periodic_map.get(periodic_time):
                task_periodic_map[periodic_time] = [task]
            else:
                task_periodic_map[periodic_time].append(task)
        for periodic_time, tasks in task_periodic_map.items():
            task_count = len(tasks)
            days = math.ceil(periodic_time / (24 * 3600))
            if days < 2:
                continue
            # 根据周期与任务数量分布任务执行时间
            task_step = math.ceil(task_count / days)
            for i in range(days):
                _tasks = tasks[i * task_step: (i + 1) * task_step]
                for task in _tasks:
                    self._add_date(task, i)

    def _request(self, task_list):
        url = f"{SCANNER_SERVER}/v1/tasks:batchModify"
        resp = requests.post(url, json={"custom_tasks": task_list})
        if resp.ok and resp.json()["code"] == 0:
            task_ids = [t["task_id"] for t in task_list]
            self.collection_bk.update(
                {"task_id": {"$in": task_ids}},
                {"$set": {"update_to_db": True}},
                multi=True
            )
            print(f"success  {task_ids}")

    def _update_start_at(self):
        """通过接口修改任务执行时间"""

        tasks = self.collection_bk.find({"update_to_db": False}, {"update_to_db": 0})
        task_list = []
        for task_body in tasks:
            _id = task_body.pop("_id", "")
            start_at = task_body.pop("next_start_at", "")
            task_body["task_config"]["start_at"] = start_at
            task_list.append(task_body)
            if len(task_list) == 30:
                self._request(task_list)
                task_list = []
        # 最后可能不够30个
        if task_list:
            self._request(task_list)

    def _save_start_at_to_db(self):
        """将理想执行时间数据存储到数据库"""
        db_data = {t["task_id"]: 1 for t in self.collection_bk.find({}, {"task_id": 1})}
        for task_id, config in self.task_config_map.items():
            if db_data.get(task_id):
                continue
            start_at = arrow.get(config.get("next_start_at")).for_json()
            task_body = self.task_data_map.get(task_id)
            task_body["next_start_at"] = start_at
            task_body["update_to_db"] = False
            self.collection_bk.insert(task_body)

    def update_next_time(self):
        """计算分布时间, 必须先执行update_day， 在执行update_time"""
        # 周期天数内分布任务
        self._update_day()
        # 一天内分布任务
        self._update_time()

    def statistics(self, new=False):
        """统计当前任务执行计划分布"""
        if new:
            self.update_next_time()
        time_map = {}
        now = datetime.datetime.utcnow().date()
        days = 30
        after_date = now + datetime.timedelta(days=days)
        month_map = {f"{(now + datetime.timedelta(days=i)).strftime('%Y-%m-%d')}": 0 for i in range(days)}
        month_hour_map = {}
        for task_id, config in self.task_config_map.items():
            next_at = config.get("next_start_at")
            periodic_time = config.get("periodic_time")
            if periodic_time not in time_map:
                time_map[periodic_time] = 1
            else:
                time_map[periodic_time] += 1
            # 统计最近 days 天任务执行分布
            while now <= next_at.date() < after_date:
                month_map[next_at.strftime('%Y-%m-%d')] += 1
                if next_at.strftime('%Y-%m-%d') not in month_hour_map:
                    month_hour_map[next_at.strftime('%Y-%m-%d')] = [next_at]
                else:
                    month_hour_map[next_at.strftime('%Y-%m-%d')].append(next_at)
                next_at = next_at + datetime.timedelta(seconds=periodic_time)
        day_hour_map = {}
        # 统计每天任务分布
        for date_str, time_list in month_hour_map.items():
            item = {}
            for t in time_list:
                hour = t.strftime('%H')
                if hour not in item:
                    item[hour] = 1
                else:
                    item[hour] += 1
            day_hour_map.update({date_str: item})
        return {"time_map": time_map, "month_map": month_map, "day_hour_map": day_hour_map}

    def decentralize(self):
        """执行任务分布"""
        self.update_next_time()
        # 把数据存在数据库
        self._save_start_at_to_db()
        # 调用接口修改任务配置
        self._update_start_at()


if __name__ == "__main__":
    DecentralizeTask().statistics()
